package uk.co.jgo2rdf.uniprotidmapping;

import com.csvreader.CsvReader;
import com.sun.xml.internal.bind.v2.model.core.ID;
import org.apache.commons.net.ftp.FTP;
import org.apache.commons.net.ftp.FTPClient;
import org.apache.log4j.Logger;
import org.apache.lucene.analysis.KeywordAnalyzer;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.Field;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.index.IndexWriter;
import org.apache.lucene.index.IndexWriterConfig;
import org.apache.lucene.index.Term;
import org.apache.lucene.search.*;
import org.apache.lucene.store.NIOFSDirectory;
import org.apache.lucene.util.Version;

import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.net.MalformedURLException;
import java.net.URL;
import java.nio.charset.Charset;
import java.util.Arrays;
import java.util.Collection;
import java.util.HashSet;
import java.util.Set;
import java.util.zip.GZIPInputStream;

/**
 * Created by IntelliJ IDEA.
 * User: mhindle
 * Date: 07/03/12
 * Time: 14:43
 * To change this template use File | Settings | File Templates.
 */
public class UniprotIdMappingIndex {

    private static String ID = "ID";
    private static String UNIPROT = "UNIPROT";
    private static String DB = "DB";


    private Set<String> databaseFrom = new HashSet<String>(); //stores all database ids
    private NIOFSDirectory directory;
    private String indexDir;

    private IndexSearcher indexSearcher; //is automatically init after parsing

    private static final Logger log = Logger.getLogger(UniprotIdMappingIndex.class);
    private static String url = "ftp://ftp.ebi.ac.uk/pub/databases/uniprot/current_release/knowledgebase/idmapping/idmapping.dat.gz";


    /**
     * @param indexDir  directory of lucene index
     * @param loadIndex if true it will load the existing index at the dir else it is deleted
     */
    public UniprotIdMappingIndex(String indexDir, boolean loadIndex) throws IOException {
        this.indexDir = indexDir;

        File f = new File(indexDir);
        if (!loadIndex) { //this is the fist run
            f.delete();
            f.mkdirs();
            f.mkdir();
        }

        directory = new NIOFSDirectory(f);

        if (loadIndex) {
            indexSearcher = new IndexSearcher(IndexReader.open(directory, true), null);
        }
    }

    public Set<String> getDatabasesFrom() {
        return databaseFrom;
    }

    public void parse() throws IOException {
        parse(new URL(url));
    }

    public void parse(URL url) throws IOException {
        if (url.getProtocol().equalsIgnoreCase("ftp")) {
            log.info("Fetch " + url.getFile() + " from " + url.getHost());

            FTPClient ftp = new FTPClient();
            try {
                ftp.connect(url.getHost());
                ftp.login("anonymous", "anonymous");
                ftp.setFileTransferMode(FTP.STREAM_TRANSFER_MODE);
                ftp.setFileType(FTP.BINARY_FILE_TYPE);
                if (url.getFile().endsWith(".gz")) {
                    parse(new GZIPInputStream(
                            ftp.retrieveFileStream(url.getFile())));
                } else {
                    parse(ftp.retrieveFileStream(url.getFile()));
                }
            } finally {
                ftp.disconnect();
            }
        } else {
            if (url.getFile().endsWith(".gz")) {
                parse(new GZIPInputStream(url.openStream()));
            } else {
                parse(url.openStream());
            }
        }
    }

    public void parse(InputStream is) throws IOException {

        if (indexSearcher != null)
            indexSearcher.close();

        IndexWriterConfig config = new IndexWriterConfig(Version.LUCENE_30,
                new KeywordAnalyzer());
        IndexWriter writer = new IndexWriter(directory, config);

        CsvReader csvReader = new CsvReader(is,
                '\t', Charset.defaultCharset());
        csvReader.setSafetySwitch(false);
        try {
            csvReader.setComment('!');
            csvReader.setUseComments(true);
            csvReader.setTrimWhitespace(false);
            csvReader.setHeaders(new String[]{UNIPROT, DB, ID});

            int i = 0;
            while (csvReader.readRecord()) {
                if (i % 10000 == 0)
                    System.out.println("Processed " + i + " valid records");

                String uniprot = csvReader.get(UNIPROT);
                if (uniprot == null || uniprot.length() == 0) continue;

                String id = csvReader.get(ID);
                if (id == null || id.length() == 0) continue;

                String db = csvReader.get(DB);
                if (db == null || db.length() == 0) continue;

                i++;

                databaseFrom.add(db.trim().toUpperCase());
                Document doc = new Document();
                doc.add(new Field(UNIPROT,
                        uniprot.trim().toUpperCase(),
                        Field.Store.YES, Field.Index.ANALYZED_NO_NORMS));
                doc.add(new Field(ID,
                        id.trim().toUpperCase(),
                        Field.Store.YES, Field.Index.ANALYZED_NO_NORMS));
                doc.add(new Field(DB,
                        db.trim().toUpperCase(),
                        Field.Store.YES, Field.Index.ANALYZED_NO_NORMS));
                writer.addDocument(doc);

                int version = id.lastIndexOf('.');
                if (version > 0) {
                    Document docuv = new Document();
                    docuv.add(new Field(UNIPROT,
                            uniprot.trim().toUpperCase(),
                            Field.Store.YES, Field.Index.ANALYZED_NO_NORMS));
                    docuv.add(new Field(ID,
                            id.substring(0, version).trim().toUpperCase(),
                            Field.Store.YES, Field.Index.ANALYZED_NO_NORMS));
                    docuv.add(new Field(DB,
                            db.trim().toUpperCase(),
                            Field.Store.YES, Field.Index.ANALYZED_NO_NORMS));
                    writer.addDocument(docuv);
                }
            }
        } finally {
            csvReader.close();
            is.close();
            writer.prepareCommit();
            writer.commit();
            writer.close(true);
        }
        indexSearcher = new IndexSearcher(IndexReader.open(directory, true), null);
    }

    /**
     * @param accession      the query accession
     * @param db the target database to translate into
     * @return accessions from the target database
     * @throws IOException
     */
    public Set<String> getAccession(String accession, String db) throws IOException {
        final Set<String> accessions = new HashSet<String>();

        BooleanQuery bq = new BooleanQuery();
        bq.add(new TermQuery(new Term(ID, accession.toUpperCase())), BooleanClause.Occur.MUST);
        if (db != null)
            bq.add(new TermQuery(new Term(DB, db.toUpperCase())), BooleanClause.Occur.MUST);

        final Collector docCollector = new Collector() {

            private int docBase;

            @Override
            public void setScorer(Scorer scorer) throws IOException {
            }

            @Override
            public void collect(int doc) throws IOException {
                try {
                    accessions.add(indexSearcher.doc(doc + docBase).get(UNIPROT));
                } catch (Exception e) {
                    log.error(e);
                    e.printStackTrace();
                }
            }

            @Override
            public void setNextReader(IndexReader reader, int docBase) throws IOException {
                this.docBase = docBase;
            }

            @Override
            public boolean acceptsDocsOutOfOrder() {
                return true;
            }
        };

        indexSearcher.search(bq, docCollector);

        return accessions;
    }

    public static void main(String[] args) throws IOException {
       //URL url = new URL("file:///home/mhindle/UniProt/idmapping.dat");
       UniprotIdMappingIndex index = new UniprotIdMappingIndex("/home/mhindle/UniProt/idmapping", true);
       //index.parse(url);

       Set<String> uniprots = index.getAccession("At3g17750", "TAIR");
       System.out.println(Arrays.toString(uniprots.toArray(new String[uniprots.size()])));
       Set<String>  uniprots2 = index.getAccession("At3g17750",null);
       System.out.println(Arrays.toString(uniprots2.toArray(new String[uniprots2.size()])));
    }

}
